home *** CD-ROM | disk | FTP | other *** search
/ EuroCD 3 / EuroCD 3.iso / Programming / Python-1.4 / Demo / classes / Complex.py < prev    next >
Text File  |  1998-06-24  |  8KB  |  299 lines
  1. # Complex numbers
  2. # ---------------
  3.  
  4. # [Now that Python has a complex data type built-in, this is not very
  5. # useful, but it's still a nice example class]
  6.  
  7. # This module represents complex numbers as instances of the class Complex.
  8. # A Complex instance z has two data attribues, z.re (the real part) and z.im
  9. # (the imaginary part).  In fact, z.re and z.im can have any value -- all
  10. # arithmetic operators work regardless of the type of z.re and z.im (as long
  11. # as they support numerical operations).
  12. #
  13. # The following functions exist (Complex is actually a class):
  14. # Complex([re [,im]) -> creates a complex number from a real and an imaginary part
  15. # IsComplex(z) -> true iff z is a complex number (== has .re and .im attributes)
  16. # ToComplex(z) -> a complex number equal to z; z itself if IsComplex(z) is true
  17. #                 if z is a tuple(re, im) it will also be converted
  18. # PolarToComplex([r [,phi [,fullcircle]]]) ->
  19. #    the complex number z for which r == z.radius() and phi == z.angle(fullcircle)
  20. #    (r and phi default to 0)
  21. # exp(z) -> returns the complex exponential of z. Equivalent to pow(math.e,z).
  22. #
  23. # Complex numbers have the following methods:
  24. # z.abs() -> absolute value of z
  25. # z.radius() == z.abs()
  26. # z.angle([fullcircle]) -> angle from positive X axis; fullcircle gives units
  27. # z.phi([fullcircle]) == z.angle(fullcircle)
  28. #
  29. # These standard functions and unary operators accept complex arguments:
  30. # abs(z)
  31. # -z
  32. # +z
  33. # not z
  34. # repr(z) == `z`
  35. # str(z)
  36. # hash(z) -> a combination of hash(z.re) and hash(z.im) such that if z.im is zero
  37. #            the result equals hash(z.re)
  38. # Note that hex(z) and oct(z) are not defined.
  39. #
  40. # These conversions accept complex arguments only if their imaginary part is zero:
  41. # int(z)
  42. # long(z)
  43. # float(z)
  44. #
  45. # The following operators accept two complex numbers, or one complex number
  46. # and one real number (int, long or float):
  47. # z1 + z2
  48. # z1 - z2
  49. # z1 * z2
  50. # z1 / z2
  51. # pow(z1, z2)
  52. # cmp(z1, z2)
  53. # Note that z1 % z2 and divmod(z1, z2) are not defined,
  54. # nor are shift and mask operations.
  55. #
  56. # The standard module math does not support complex numbers.
  57. # (I suppose it would be easy to implement a cmath module.)
  58. #
  59. # Idea:
  60. # add a class Polar(r, phi) and mixed-mode arithmetic which
  61. # chooses the most appropriate type for the result:
  62. # Complex for +,-,cmp
  63. # Polar   for *,/,pow
  64.  
  65.  
  66. import types, math
  67.  
  68. twopi = math.pi*2.0
  69. halfpi = math.pi/2.0
  70.  
  71. def IsComplex(obj):
  72.     return hasattr(obj, 're') and hasattr(obj, 'im')
  73.  
  74. def ToComplex(obj):
  75.     if IsComplex(obj):
  76.         return obj
  77.     elif type(obj) == types.TupleType:
  78.         return apply(Complex, obj)
  79.     else:
  80.         return Complex(obj)
  81.  
  82. def PolarToComplex(r = 0, phi = 0, fullcircle = twopi):
  83.     phi = phi * (twopi / fullcircle)
  84.     return Complex(math.cos(phi)*r, math.sin(phi)*r)
  85.  
  86. def Re(obj):
  87.     if IsComplex(obj):
  88.         return obj.re
  89.     else:
  90.         return obj
  91.  
  92. def Im(obj):
  93.     if IsComplex(obj):
  94.         return obj.im
  95.     else:
  96.         return obj
  97.  
  98. class Complex:
  99.  
  100.     def __init__(self, re=0, im=0):
  101.         if IsComplex(re):
  102.             im = i + Complex(0, re.im)
  103.             re = re.re
  104.         if IsComplex(im):
  105.             re = re - im.im
  106.             im = im.re
  107.         self.__dict__['re'] = re
  108.         self.__dict__['im'] = im
  109.     
  110.     def __setattr__(self, name, value):
  111.             raise TypeError, 'Complex numbers are immutable'
  112.  
  113.     def __hash__(self):
  114.         if not self.im: return hash(self.re)
  115.         mod = sys.maxint + 1L
  116.         return int((hash(self.re) + 2L*hash(self.im) + mod) % (2L*mod) - mod)
  117.  
  118.     def __repr__(self):
  119.         if not self.im:
  120.             return 'Complex(%s)' % `self.re`
  121.         else:
  122.             return 'Complex(%s, %s)' % (`self.re`, `self.im`)
  123.  
  124.     def __str__(self):
  125.         if not self.im:
  126.             return `self.re`
  127.         else:
  128.             return 'Complex(%s, %s)' % (`self.re`, `self.im`)
  129.  
  130.     def __neg__(self):
  131.         return Complex(-self.re, -self.im)
  132.  
  133.     def __pos__(self):
  134.         return self
  135.  
  136.     def __abs__(self):
  137.         # XXX could be done differently to avoid overflow!
  138.         return math.sqrt(self.re*self.re + self.im*self.im)
  139.  
  140.     def __int__(self):
  141.         if self.im:
  142.             raise ValueError, "can't convert Complex with nonzero im to int"
  143.         return int(self.re)
  144.  
  145.     def __long__(self):
  146.         if self.im:
  147.             raise ValueError, "can't convert Complex with nonzero im to long"
  148.         return long(self.re)
  149.  
  150.     def __float__(self):
  151.         if self.im:
  152.             raise ValueError, "can't convert Complex with nonzero im to float"
  153.         return float(self.re)
  154.  
  155.     def __cmp__(self, other):
  156.         other = ToComplex(other)
  157.         return cmp((self.re, self.im), (other.re, other.im))
  158.  
  159.     def __rcmp__(self, other):
  160.         other = ToComplex(other)
  161.         return cmp(other, self)
  162.     
  163.     def __nonzero__(self):
  164.         return not (self.re == self.im == 0)
  165.  
  166.     abs = radius = __abs__
  167.  
  168.     def angle(self, fullcircle = twopi):
  169.         return (fullcircle/twopi) * ((halfpi - math.atan2(self.re, self.im)) % twopi)
  170.  
  171.     phi = angle
  172.  
  173.     def __add__(self, other):
  174.         other = ToComplex(other)
  175.         return Complex(self.re + other.re, self.im + other.im)
  176.  
  177.     __radd__ = __add__
  178.  
  179.     def __sub__(self, other):
  180.         other = ToComplex(other)
  181.         return Complex(self.re - other.re, self.im - other.im)
  182.  
  183.     def __rsub__(self, other):
  184.         other = ToComplex(other)
  185.         return other - self
  186.  
  187.     def __mul__(self, other):
  188.         other = ToComplex(other)
  189.         return Complex(self.re*other.re - self.im*other.im,
  190.                        self.re*other.im + self.im*other.re)
  191.  
  192.     __rmul__ = __mul__
  193.  
  194.     def __div__(self, other):
  195.         other = ToComplex(other)
  196.         d = float(other.re*other.re + other.im*other.im)
  197.         if not d: raise ZeroDivisionError, 'Complex division'
  198.         return Complex((self.re*other.re + self.im*other.im) / d,
  199.                        (self.im*other.re - self.re*other.im) / d)
  200.  
  201.     def __rdiv__(self, other):
  202.         other = ToComplex(other)
  203.         return other / self
  204.  
  205.     def __pow__(self, n, z=None):
  206.         if z is not None:
  207.             raise TypeError, 'Complex does not support ternary pow()'
  208.         if IsComplex(n):
  209.             if n.im: 
  210.               if self.im: raise TypeError, 'Complex to the Complex power'
  211.               else: return exp(math.log(self.re)*n)
  212.             n = n.re
  213.         r = pow(self.abs(), n)
  214.         phi = n*self.angle()
  215.         return Complex(math.cos(phi)*r, math.sin(phi)*r)
  216.     
  217.     def __rpow__(self, base):
  218.         base = ToComplex(base)
  219.         return pow(base, self)
  220.         
  221. def exp(z):
  222.     r = math.exp(z.re)
  223.     return Complex(math.cos(z.im)*r,math.sin(z.im)*r)
  224.  
  225.  
  226. def checkop(expr, a, b, value, fuzz = 1e-6):
  227.     import sys
  228.     print '       ', a, 'and', b,
  229.     try:
  230.         result = eval(expr)
  231.     except:
  232.         result = sys.exc_type
  233.     print '->', result
  234.     if (type(result) == type('') or type(value) == type('')):
  235.         ok = result == value
  236.     else:
  237.         ok = abs(result - value) <= fuzz
  238.     if not ok:
  239.         print '!!\t!!\t!! should be', value, 'diff', abs(result - value)
  240.  
  241.  
  242. def test():
  243.     testsuite = {
  244.         'a+b': [
  245.             (1, 10, 11),
  246.             (1, Complex(0,10), Complex(1,10)),
  247.             (Complex(0,10), 1, Complex(1,10)),
  248.             (Complex(0,10), Complex(1), Complex(1,10)),
  249.             (Complex(1), Complex(0,10), Complex(1,10)),
  250.         ],
  251.         'a-b': [
  252.             (1, 10, -9),
  253.             (1, Complex(0,10), Complex(1,-10)),
  254.             (Complex(0,10), 1, Complex(-1,10)),
  255.             (Complex(0,10), Complex(1), Complex(-1,10)),
  256.             (Complex(1), Complex(0,10), Complex(1,-10)),
  257.         ],
  258.         'a*b': [
  259.             (1, 10, 10),
  260.             (1, Complex(0,10), Complex(0, 10)),
  261.             (Complex(0,10), 1, Complex(0,10)),
  262.             (Complex(0,10), Complex(1), Complex(0,10)),
  263.             (Complex(1), Complex(0,10), Complex(0,10)),
  264.         ],
  265.         'a/b': [
  266.             (1., 10, 0.1),
  267.             (1, Complex(0,10), Complex(0, -0.1)),
  268.             (Complex(0, 10), 1, Complex(0, 10)),
  269.             (Complex(0, 10), Complex(1), Complex(0, 10)),
  270.             (Complex(1), Complex(0,10), Complex(0, -0.1)),
  271.         ],
  272.         'pow(a,b)': [
  273.             (1, 10, 1),
  274.             (1, Complex(0,10), 'TypeError'),
  275.             (Complex(0,10), 1, Complex(0,10)),
  276.             (Complex(0,10), Complex(1), Complex(0,10)),
  277.             (Complex(1), Complex(0,10), 'TypeError'),
  278.             (2, Complex(4,0), 16),
  279.         ],
  280.         'cmp(a,b)': [
  281.             (1, 10, -1),
  282.             (1, Complex(0,10), 1),
  283.             (Complex(0,10), 1, -1),
  284.             (Complex(0,10), Complex(1), -1),
  285.             (Complex(1), Complex(0,10), 1),
  286.         ],
  287.     }
  288.     exprs = testsuite.keys()
  289.     exprs.sort()
  290.     for expr in exprs:
  291.         print expr + ':'
  292.         t = (expr,)
  293.         for item in testsuite[expr]:
  294.             apply(checkop, t+item)
  295.     
  296.  
  297. if __name__ == '__main__':
  298.     test()
  299.